import numpy as np

def viterbi(A, B, pi, O):
    '''
    :param A: 状态转移概率矩阵：一个状态转移到另一个状态的概率
    :param B: 观测概率矩阵：行表示某个状态，列表示某个观测，Bij表示在这个状态下产生这个观测的概率,
    设第一列观测序号为0，第二列观测的序号为1，...
    :param pi: 初始状态概率向量：以某个状态开始的概率,pi的长度即为可选的状态的个数
    :param O: 观测序列，由0，1，2...组成，O的长度即为观测次数
    :return: 最优路径I:一个数组，表示每个时刻最可能选择的状态的编号i，为0,1,2...
    '''

    #第一步：初始化:变量
    # k表示观测次数
    k = O.shape[0]
    # n表示可选状态的数目
    n = pi.shape[0]
    # node记录时刻t时，状态为i的所有路径中概率最大的路径的上一个结点，即为上一个状态的编号0，1，2...，大小为 k * n,每个元素初始化为0
    node = [[0 for i in range(n)] for j in range(k)]
    #delta记录时刻t时，状态i观测到O[t]的概率，大小为k * n，每个元素初始化为0
    delta = [[0 for i in range(n)] for j in range(k)]

    #第二部：单独初始化时刻1，时刻1的观测为O[0]，
    for i in range(n):
        delta[0][i] = pi[i] * B[i][O[0]]

    #第三步：递推计算:时刻2，3，4，....
    for t in range(1, k):
        for i in range(n):
            # 计算t时刻，状态i，在上一个时刻各个状态下得到的观测的基础上，该时刻i状态得到新观测值的最大概率
            # 即计算delta[t][i]
            # 计算node[t][i]，记录的是前一个状态j
            prob_max = 0
            j_max = 0
            for j in range(n):
                p = delta[t-1][j] * A[j][i]
                if p > prob_max:
                    prob_max = p
                    j_max = j
            delta[t][i] = prob_max * B[i][O[t]]
            node[t][i] = j_max

    #第四步：找到最优路径的终点
    #I为返回结果，是最优路径序列
    I = []
    i = np.argmax(delta[k-1])
    I.append(i)

    #第五步：由最优路径的终点回溯找到其他结点
    for t in range(k-2, -1, -1):
        I.insert(0, node[t+1][i])
        i = node[t+1][i]

    return I





if __name__ == "__main__":
    A = np.array([[0, 1],
                  [0.5, 0.5]])
    B = np.array([[0.5, 0.5],
                  [0, 1]])
    pi = np.array([0.5, 0.5])
    O = np.array([0, 1, 1, 1, 0])
    I = viterbi(A, B, pi, O)
    print(I)